package com.runjian.rundo.bmp.kernel.common.spring.context.utils;

import java.util.Map;
import java.util.Objects;
import java.util.concurrent.Callable;

import org.slf4j.MDC;

import cn.hutool.core.util.IdUtil;

/**
 *
 * @author CaoQingHua
 */
public class MdcUtils {
    public MdcUtils() {}

    public static String getTraceId() {
        return MDC.get("traceId");
    }

    public static void setTraceId(String traceId) {
        MDC.put("traceId", traceId);
    }

    public static void setTraceIdIfAbsent(String traceId) {
        if (Objects.isNull(traceId)) {
            setTraceIdIfAbsent();
        } else {
            setTraceId(traceId);
        }

    }

    public static void setTraceIdIfAbsent() {
        if (MDC.get("traceId") == null) {
            MDC.put("traceId", IdUtil.fastSimpleUUID());
        }

    }

    public static <T> Callable<T> wrap(Callable<T> callable, Map<String, String> context) {
        return () -> {
            if (context == null) {
                MDC.clear();
            } else {
                MDC.setContextMap(context);
            }

            setTraceIdIfAbsent();

            T var2;
            try {
                var2 = callable.call();
            } finally {
                MDC.clear();
            }

            return var2;
        };
    }

    public static Runnable wrap(Runnable runnable, Map<String, String> context) {
        return () -> {
            if (context == null) {
                MDC.clear();
            } else {
                MDC.setContextMap(context);
            }

            setTraceIdIfAbsent();

            try {
                runnable.run();
            } finally {
                MDC.clear();
            }

        };
    }

    public static void cleanTraceId() {
        MDC.clear();
    }
}
